#!/usr/bin/env python
# encoding: UTF-8

"""Usage: %s database

Prints a csv summary of contacts in the set."""

from db_cleanup import get_tables
import sqlite3
import sys

def main(argv):
    sum_contacts = {}
    try:
        tables = get_tables(argv[1])
    except IndexError:
        print __doc__ % argv[0]
        sys.exit(1)

    conn = sqlite3.connect(argv[1])
    conn.row_factory = sqlite3.Row
    cur = conn.cursor()

    nucls = set(("A", "G", "C", "T", "sugar", "phosphate"))
    aminos = set(("GLY", "ALA", "VAL", "LEU", "ILE", "ASP", "ASN", "GLU",
                 "GLN", "ARG", "LYS", "HIS", "PHE", "SER", "THR", "TYR",
                 "TRP", "MET", "CYS", "PRO"))
    excl = set()
    for (contacts, atoms) in tables:
        pdb = contacts.split("_")[1]
        
        cur.execute("""SELECT COUNT(*), type, directness, d.amino AS damino, a.amino AS aamino, d.atom AS datom, a.atom AS aatom
                            FROM %(contacts)s c, %(atoms)s a, %(atoms)s d
                        WHERE
                            donor_id = d.id
                            AND acceptor_id = a.id
                        GROUP BY a.amino, d.amino, type, directness"""
                        % {"contacts" : contacts, "atoms": atoms})
        for row in cur:
            if row["damino"] in nucls:
                amino = row["aamino"]
                nucl  = row["damino"]
                nucatom = row["datom"]
            else: #aamino is DNA
                amino = row["damino"]
                nucl  = row["aamino"]
                nucatom = row["aatom"]

            if amino not in aminos:
                excl.add(amino)
                continue

            if nucatom.endswith("*"):
                nucl = "sugar"
            elif nucatom in ("P", "O1P", "O2P"):
                nucl = "phosphate"

                
            type = row["type"]
            directness = row["directness"]
            
            sum_contacts[amino] = sum_contacts.get(amino, {})
            sum_contacts[amino][nucl] = sum_contacts[amino].get(nucl, {})
            sum_contacts[amino][nucl][type] = sum_contacts[amino][nucl].get(type, {})
            sum_contacts[amino][nucl][type][directness] = sum_contacts[amino][nucl][type].get(directness,0) + row[0]

    print "Hydrogen bonds direct:"
    print_contacts(sum_contacts, nucls, "H", "D")
    print
    print "Hydrogen bonds water-mediated:"
    print_contacts(sum_contacts, nucls, "H", "WX")
    print
    print "Near contacts direct:"
    print_contacts(sum_contacts, nucls, "N", "D")
    print
    print "Near contacts water-mediated:"
    print_contacts(sum_contacts, nucls, "N", "WX")
    print
    print "Excluded:", excl
    
def print_contacts(contacts, nucls, type="H", directness="D"):
    print "   ,",
    nucls = list(nucls)
    nucls.sort()
    print ",".join(nucls)
    for amino in sorted(contacts):
        print "%3s," % amino,
        nums = []
        for nucl in nucls:
            cs = contacts[amino].get(nucl, {}).get(type, {})
            n = cs.get(directness, 0) if len(directness) == 1 else sum(cs.get(d, 0) for d in directness)
            nums.append(str(n))
        print ", ".join(nums)

if __name__ == "__main__":
    main(sys.argv)
